{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-12-02T11:36:16.413716Z",
     "start_time": "2019-12-02T11:36:16.404566Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "# pylint: disable=missing-docstring\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "from __future__ import unicode_literals\n",
    "\n",
    "import sys\n",
    "sys.path.append('../cleverhans/')\n",
    "import warnings\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch import optim\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-12-02T11:36:19.515700Z",
     "start_time": "2019-12-02T11:36:19.482924Z"
    }
   },
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'tensorflow' has no attribute 'GraphKeys'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-0ec80f32bf89>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattacks\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mFastGradientMethod\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/anaconda3/envs/dl/lib/python3.7/site-packages/cleverhans/attacks/__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mutils\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattacks_tf\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mSPSAAdam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmargin_logit_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mTensorAdam\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mModel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mCallableModelWrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mwrapper_warning\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwrapper_warning_logits\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/dl/lib/python3.7/site-packages/cleverhans/attacks_tf.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompat\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mreduce_any\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompat\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0msoftmax_cross_entropy_with_logits\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mutils_tf\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mutils\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/dl/lib/python3.7/site-packages/cleverhans/utils_tf.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    339\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    340\u001b[0m def kl_with_logits(p_logits, q_logits, scope=None,\n\u001b[0;32m--> 341\u001b[0;31m                    loss_collection=tf.GraphKeys.REGULARIZATION_LOSSES):\n\u001b[0m\u001b[1;32m    342\u001b[0m   \"\"\"Helper function to compute kl-divergence KL(p || q)\n\u001b[1;32m    343\u001b[0m   \"\"\"\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'tensorflow' has no attribute 'GraphKeys'"
     ]
    }
   ],
   "source": [
    "from cleverhans.attacks import FastGradientMethod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-12-02T11:34:39.628908Z",
     "start_time": "2019-12-02T11:34:39.606847Z"
    }
   },
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'flags' from 'cleverhans.compat' (/home/shenchenkai/anaconda3/envs/dl/lib/python3.7/site-packages/cleverhans/compat.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-d433019ba447>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompat\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mflags\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mCallableModelWrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mAccuracyReport\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mcleverhans\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils_pytorch\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mconvert_pytorch_model_to_tf\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'flags' from 'cleverhans.compat' (/home/shenchenkai/anaconda3/envs/dl/lib/python3.7/site-packages/cleverhans/compat.py)"
     ]
    }
   ],
   "source": [
    "from cleverhans.compat import flags\n",
    "from cleverhans.model import CallableModelWrapper\n",
    "from cleverhans.utils import AccuracyReport\n",
    "from cleverhans.utils_pytorch import convert_pytorch_model_to_tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FLAGS = flags.FLAGS\n",
    "\n",
    "NB_EPOCHS = 6\n",
    "BATCH_SIZE = 128\n",
    "LEARNING_RATE = .001\n",
    "\n",
    "warnings.warn(\"convert_pytorch_model_to_tf is deprecated, switch to\"\n",
    "              + \" dedicated PyTorch support provided by CleverHans v4.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PytorchMnistModel(nn.Module):\n",
    "  \"\"\" Basic MNIST model from github\n",
    "  https://github.com/rickiepark/pytorch-examples/blob/master/mnist.ipynb\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self):\n",
    "    super(PytorchMnistModel, self).__init__()\n",
    "    # input is 28x28\n",
    "    # padding=2 for same padding\n",
    "    self.conv1 = nn.Conv2d(1, 32, 5, padding=2)\n",
    "    # feature map size is 14*14 by pooling\n",
    "    # padding=2 for same padding\n",
    "    self.conv2 = nn.Conv2d(32, 64, 5, padding=2)\n",
    "    # feature map size is 7*7 by pooling\n",
    "    self.fc1 = nn.Linear(64 * 7 * 7, 1024)\n",
    "    self.fc2 = nn.Linear(1024, 10)\n",
    "\n",
    "  def forward(self, x):\n",
    "    x = F.max_pool2d(F.relu(self.conv1(x)), 2)\n",
    "    x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n",
    "    x = x.view(-1, 64 * 7 * 7)  # reshape Variable\n",
    "    x = F.relu(self.fc1(x))\n",
    "    x = self.fc2(x)\n",
    "    return F.log_softmax(x, dim=-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist_tutorial(nb_epochs=NB_EPOCHS, batch_size=BATCH_SIZE,\n",
    "                   train_end=-1, test_end=-1, learning_rate=LEARNING_RATE):\n",
    "  \"\"\"\n",
    "  MNIST cleverhans tutorial\n",
    "  :param nb_epochs: number of epochs to train model\n",
    "  :param batch_size: size of training batches\n",
    "  :param learning_rate: learning rate for training\n",
    "  :return: an AccuracyReport object\n",
    "  \"\"\"\n",
    "  # Train a pytorch MNIST model\n",
    "  torch_model = PytorchMnistModel()\n",
    "  if torch.cuda.is_available():\n",
    "    torch_model = torch_model.cuda()\n",
    "  report = AccuracyReport()\n",
    "\n",
    "  train_loader = torch.utils.data.DataLoader(\n",
    "      datasets.MNIST('data', train=True, download=True,\n",
    "                     transform=transforms.ToTensor()),\n",
    "      batch_size=batch_size, shuffle=True)\n",
    "  test_loader = torch.utils.data.DataLoader(\n",
    "      datasets.MNIST('data', train=False, transform=transforms.ToTensor()),\n",
    "      batch_size=batch_size)\n",
    "\n",
    "  # Truncate the datasets so that our test run more quickly\n",
    "  train_loader.dataset.train_data = train_loader.dataset.train_data[\n",
    "      :train_end]\n",
    "  test_loader.dataset.test_data = test_loader.dataset.test_data[:test_end]\n",
    "\n",
    "  # Train our model\n",
    "  optimizer = optim.Adam(torch_model.parameters(), lr=learning_rate)\n",
    "  train_loss = []\n",
    "\n",
    "  total = 0\n",
    "  correct = 0\n",
    "  step = 0\n",
    "  for _epoch in range(nb_epochs):\n",
    "    for xs, ys in train_loader:\n",
    "      xs, ys = Variable(xs), Variable(ys)\n",
    "      if torch.cuda.is_available():\n",
    "        xs, ys = xs.cuda(), ys.cuda()\n",
    "      optimizer.zero_grad()\n",
    "      preds = torch_model(xs)\n",
    "      loss = F.nll_loss(preds, ys)\n",
    "      loss.backward()  # calc gradients\n",
    "      train_loss.append(loss.data.item())\n",
    "      optimizer.step()  # update gradients\n",
    "\n",
    "      preds_np = preds.cpu().detach().numpy()\n",
    "      correct += (np.argmax(preds_np, axis=1) == ys.cpu().detach().numpy()).sum()\n",
    "      total += train_loader.batch_size\n",
    "      step += 1\n",
    "      if total % 1000 == 0:\n",
    "        acc = float(correct) / total\n",
    "        print('[%s] Training accuracy: %.2f%%' % (step, acc * 100))\n",
    "        total = 0\n",
    "        correct = 0\n",
    "\n",
    "  # Evaluate on clean data\n",
    "  total = 0\n",
    "  correct = 0\n",
    "  for xs, ys in test_loader:\n",
    "    xs, ys = Variable(xs), Variable(ys)\n",
    "    if torch.cuda.is_available():\n",
    "      xs, ys = xs.cuda(), ys.cuda()\n",
    "\n",
    "    preds = torch_model(xs)\n",
    "    preds_np = preds.cpu().detach().numpy()\n",
    "\n",
    "    correct += (np.argmax(preds_np, axis=1) == ys.cpu().detach().numpy()).sum()\n",
    "    total += len(xs)\n",
    "\n",
    "  acc = float(correct) / total\n",
    "  report.clean_train_clean_eval = acc\n",
    "  print('[%s] Clean accuracy: %.2f%%' % (step, acc * 100))\n",
    "\n",
    "  # We use tf for evaluation on adversarial data\n",
    "  sess = tf.Session()\n",
    "  x_op = tf.placeholder(tf.float32, shape=(None, 1, 28, 28,))\n",
    "\n",
    "  # Convert pytorch model to a tf_model and wrap it in cleverhans\n",
    "  tf_model_fn = convert_pytorch_model_to_tf(torch_model)\n",
    "  cleverhans_model = CallableModelWrapper(tf_model_fn, output_layer='logits')\n",
    "\n",
    "  # Create an FGSM attack\n",
    "  fgsm_op = FastGradientMethod(cleverhans_model, sess=sess)\n",
    "  fgsm_params = {'eps': 0.3,\n",
    "                 'clip_min': 0.,\n",
    "                 'clip_max': 1.}\n",
    "  adv_x_op = fgsm_op.generate(x_op, **fgsm_params)\n",
    "  adv_preds_op = tf_model_fn(adv_x_op)\n",
    "\n",
    "  # Run an evaluation of our model against fgsm\n",
    "  total = 0\n",
    "  correct = 0\n",
    "  for xs, ys in test_loader:\n",
    "    adv_preds = sess.run(adv_preds_op, feed_dict={x_op: xs})\n",
    "    correct += (np.argmax(adv_preds, axis=1) == ys.cpu().detach().numpy()).sum()\n",
    "    total += test_loader.batch_size\n",
    "\n",
    "  acc = float(correct) / total\n",
    "  print('Adv accuracy: {:.3f}'.format(acc * 100))\n",
    "  report.clean_train_adv_eval = acc\n",
    "  return report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(_=None):\n",
    "  from cleverhans_tutorials import check_installation\n",
    "  check_installation(__file__)\n",
    "\n",
    "  mnist_tutorial(nb_epochs=FLAGS.nb_epochs,\n",
    "                 batch_size=FLAGS.batch_size,\n",
    "                 learning_rate=FLAGS.learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == '__main__':\n",
    "  flags.DEFINE_integer('nb_epochs', NB_EPOCHS,\n",
    "                       'Number of epochs to train model')\n",
    "  flags.DEFINE_integer('batch_size', BATCH_SIZE,\n",
    "                       'Size of training batches')\n",
    "  flags.DEFINE_float('learning_rate', LEARNING_RATE,\n",
    "                     'Learning rate for training')\n",
    "\n",
    "  tf.app.run()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
